import os
import cv2
import numpy as np

def extract_chars_from_worksheet(image_path, output_dir='chars'):
    """ 专门用于生字表汉字提取的函数 """
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 1. 图像读取与预处理
    print("1. 图像读取与预处理...")
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"无法读取图像文件：{image_path}")
    
    # 尺寸标准化（保持宽高比）
    height, width = img.shape[:2]
    if max(height, width) > 1500:
        scale = 1500 / max(height, width)
        img = cv2.resize(img, None, fx=scale, fy=scale)
    
    # 2. 针对生字表的特殊预处理
    print("2. 生字表特殊预处理...")
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 强化方格和汉字对比度
    blurred = cv2.GaussianBlur(gray, (5,5), 0)  # 保持模糊核大小
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
    enhanced = clahe.apply(blurred)
    
    # 3. 提取方格区域
    print("3. 提取方格区域...")
    binary = cv2.adaptiveThreshold(enhanced, 255, 
                                 cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                 cv2.THRESH_BINARY_INV, 45, 15)  # 保持阈值参数
    
    # 4. 检测并去除表格线
    print("4. 处理表格线...")
    # 水平线检测
    horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (50,1))  # 保持水平线检测核大小
    horizontal = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel, iterations=2)
    
    # 垂直线检测
    vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1,50))  # 保持垂直线检测核大小
    vertical = cv2.morphologyEx(binary, cv2.MORPH_OPEN, vertical_kernel, iterations=2)
    
    # 合并表格线并去除
    table_lines = cv2.add(horizontal, vertical)
    no_lines = cv2.subtract(binary, table_lines)
    
    # 5. 汉字区域增强
    print("5. 汉字区域增强...")
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))  # 保持形态学操作的核大小
    enhanced_chars = cv2.morphologyEx(no_lines, cv2.MORPH_CLOSE, kernel, iterations=2)
    
    # 6. 定位汉字方格
    print("6. 定位汉字方格...")
    contours, _ = cv2.findContours(enhanced_chars, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 过滤条件
    min_area = 1000  # 保持最小面积阈值
    max_area = 12000  # 调整最大面积阈值
    valid_cells = []
    
    for cnt in contours:
        x,y,w,h = cv2.boundingRect(cnt)
        area = w * h
        
        # 只保留符合方格特征的区域
        if min_area < area < max_area and 0.7 < w/h < 1.3:  # 放宽松高比范围
            valid_cells.append((x,y,w,h))
    
    # 7. 按行列排序方格
    print("7. 排序汉字方格...")
    if valid_cells:
        # 计算平均行高和列宽
        avg_height = np.mean([h for _,_,_,h in valid_cells])
        avg_width = np.mean([w for _,_,w,_ in valid_cells])
        
        # 按行和列排序
        valid_cells.sort(key=lambda c: (c[1]//int(avg_height*0.9), c[0]//int(avg_width*0.9)))
    
    # 8. 提取每个方格中的汉字
    print(f"8. 正在提取 {len(valid_cells)} 个汉字...")
    for i, (x,y,w,h) in enumerate(valid_cells):
        # 从原始图像提取方格区域（扩大10%范围）
        pad = int(min(w,h)*0.1)
        roi = enhanced[max(0,y-pad):min(height,y+h+pad), 
                      max(0,x-pad):min(width,x+w+pad)]
        
        # 对单个方格进行二值化
        _, cell_binary = cv2.threshold(roi, 0, 255, 
                                      cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)
        
        # 提取最大连通区域（去除小噪点）
        num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(cell_binary)
        if num_labels > 1:
            # 找到第二大的区域（第一大为背景）
            sizes = stats[1:, -1]
            max_idx = np.argmax(sizes) + 1
            char_img = np.uint8(labels == max_idx) * 255
            
            # 裁剪到字符实际区域
            x2,y2,w2,h2,_ = stats[max_idx]
            char_img = char_img[y2:y2+h2, x2:x2+w2]
            
            # 标准化输出
            char_img = cv2.copyMakeBorder(char_img, 20,20,20,20, 
                                        cv2.BORDER_CONSTANT, value=0)
            char_img = cv2.resize(char_img, (100,100), 
                                interpolation=cv2.INTER_AREA)
            
            # 保存汉字
            cv2.imwrite(f"{output_dir}/hanzi_{i+1:03d}.png", char_img)
    
    print(f"处理完成！共提取 {len(valid_cells)} 个汉字")

if __name__ == "__main__":
    input_image = "hanzi1.jpg"
    
    if not os.path.exists(input_image):
        print(f"错误：输入图像 {input_image} 不存在")
    else:
        try:
            extract_chars_from_worksheet(input_image)
        except Exception as e:
            print(f"处理过程中发生错误：{str(e)}")



